/*
 *    Copyright 2010 Vanessa Williams
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */

/*
 * BaseCount.java
 * Created on Feb 21, 2010
 */

package bioformatica.annotation;

import bioformatica.sequence.FASTASequence;
import bioformatica.sequence.Sequence;
import bioformatica.symbol.Base;
import bioformatica.util.CloseableIterator;
import org.apache.commons.cli.*;

import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.TreeMap;
import java.util.logging.Logger;

/**
 * Main class to calculate base counts, base frequencies, and GC content.
 *
 * Takes the following options on the command-line:
 * <ul>
 * <li>-seq &lt;sequence file in FASTA format&gt;</li>
 * <li>-cmd &lt;count (for base counts) | frequencey (for base frequencies) | GC (for global GC content)&gt;</li>
 * <li>-all include all symbols; if absent, only include ACGT
 * </ul>
 */
public class BaseCount {

    private static final Logger LOG = Logger.getLogger(BaseCount.class.getName());

    Map<Base,Integer> counts = new HashMap<Base,Integer>();
    int baseCount = 0;
    FASTASequence sequence;

    /**
     * Constructor
     *
     * @param fastaFilename file containing sequence in FASTA format. Must begin with header line (staring with ">")
     */
    public BaseCount(String fastaFilename) {

        CloseableIterator<Byte> bytes = null;
        try {
            this.sequence = new FASTASequence(new File(fastaFilename));
            bytes = sequence.iterator();
            byte b;
            while (bytes.hasNext()) {

                b = bytes.next();
                Base base = Base.getBase(b);

                Integer numBase = counts.get(base);
                if (numBase == null) {
                    counts.put(base, 1);
                }
                else {
                    counts.put(base, numBase+1);
                }

                baseCount++;

            }
        } catch (IOException e) {
            LOG.warning("File not found or IOException" + e.getMessage());
        }
        finally {
            if (bytes != null) bytes.close();
        }

    }

    /**
     * Get count of bases
     * @param allSymbols true for all symbols; false for just ACGT
     * @return Map of characters (bases) and integer counts
     */
    public Map<Character, Integer> getBaseCount(boolean allSymbols) {

        Map<Character, Integer> results = new TreeMap<Character, Integer>();
        if (!allSymbols) {
            results.put('A', counts.get(Base.A));
            results.put('C', counts.get(Base.C));
            results.put('G', counts.get(Base.G));
            if (sequence.getAlphabet() == Sequence.Alphabet.RNA) {
                results.put('U', counts.get(Base.U));
            }
            else {
                results.put('T', counts.get(Base.T));
            }

            if (counts.get(Base.X) != null) {
                results.put('-', counts.get(Base.X));
            }
        }
        else {
            for (Base base: counts.keySet()) {
                results.put(base.getChar(), counts.get(base));
            }

        }
        return results;
    }

    /**
     * Get base frequencies
     *
     * @param allSymbols if true, return all basses; if false, return only ACGT
     * @return
     */
    public Map<Character, Double> getBaseCountFrequency(boolean allSymbols) {
        Map<Character, Double> results = new TreeMap<Character,Double>();
        if (!allSymbols) {
            results.put('A', (double)counts.get(Base.A)/baseCount);
            results.put('C', (double)counts.get(Base.C)/baseCount);
            results.put('G', (double)counts.get(Base.G)/baseCount);
            if (sequence.getAlphabet() == Sequence.Alphabet.RNA) {
                results.put('U', (double)counts.get(Base.U)/baseCount);                
            }
            else {
                results.put('T', (double)counts.get(Base.T)/baseCount);
            }
        }
        else {
            for (Base base: counts.keySet()) {
                results.put(base.getChar(), (double)counts.get(base)/baseCount);
            }
        }
        return results;
    }

    /**
     * Get global GC content
     *
     * @return number of G + C bases in sequence
     */
    public int gcContent() {
        int cgCount = 0;
        cgCount += counts.get(Base.C);
        cgCount += counts.get(Base.G);
        return cgCount;
    }

    /**
     *
     * @param baseCounts a map of characters (base symbols) and integer counts
     * @param allSymbols if true, use all symbols; if false, use only ACGT
     * @return a string representation of the base counts
     */
    public String baseCountToString(Map<Character,Integer> baseCounts, boolean allSymbols) {
        int gapCount = 0;
        int printedBaseCount = 0;
        StringBuilder sb = new StringBuilder();
        sb.append("total bases: " + baseCount + "\n\n");
        sb.append("bases = \n\n");
        for (Character baseChar: baseCounts.keySet()) {
            if (baseChar == '-') {
                gapCount++;
                printedBaseCount += gapCount;
            }
            else {
                int count = baseCounts.get(baseChar);
                printedBaseCount += count;
                sb.append("\t" + baseChar + " : " + count + "\n");
            }

        }
        sb.append("\tGap : " + gapCount + "\n\n");
        int others = baseCount - printedBaseCount;
        if (!allSymbols) {
            sb.append("Others: " + others + "\n");
        }
        return sb.toString();
    }

    /**
     *
     * @param baseFrequencies a map of character base symbols and frequencies
     * @return a String representation of the base frequencies
     */
    public String baseFrequencyToString(Map<Character,Double> baseFrequencies) {
        int gapCount = 0;
        StringBuilder sb = new StringBuilder();
        sb.append("total bases: " + baseCount + "\n\n");
        sb.append("frequency = \n\n");
        for (Character baseChar: baseFrequencies.keySet()) {
            if (baseChar == '-') {
                gapCount++;
            }
            else {
                String freq = String.format("%.4f", baseFrequencies.get(baseChar));
                sb.append("\t" + baseChar + " : " +  freq + "\n");
            }
        }
        sb.append("\tGap : " + gapCount + "\n");
        
        return sb.toString();
    }

    public int totalBaseCount() {
        return baseCount;
    }
    
    /**
     * Main method.
     *
     * @param args
     */
    public static void main(String[] args) {

        // TODO: add an option for an output file
        
        Options options = new Options();
        options.addOption("all", false, "Show all symbols, including ambiguous bases");
        Option command = OptionBuilder.withArgName("count|freq|GC").
                                        hasArg().
                                        withDescription("Specify whether to return base count, base frequency, or GC content").
                                        create("cmd");
        Option file = OptionBuilder.withArgName("file").
                                        hasArg().
                                        withDescription("FASTA file name").
                                        create("seq");
        options.addOption(command);
        options.addOption(file);
        


        boolean allSymbols = false;
        String cmd;
        CommandLineParser parser = new GnuParser();
        try {
            // parse the command line arguments
            CommandLine line = parser.parse( options, args );
            if( line.hasOption( "all" ) ) {
                allSymbols = true;
            }
            if (line.hasOption("cmd")) {
                cmd = line.getOptionValue("cmd");
            }
            else {
                cmd = "count";
            }

            if (line.hasOption("seq")) {
                String fasta = line.getOptionValue("seq");
                BaseCount instance = new BaseCount(fasta);
                String output = null;
                if (cmd.equals("count")) {
                    Map<Character,Integer> baseCounts = instance.getBaseCount(allSymbols);
                    output = instance.baseCountToString(baseCounts, allSymbols);
                }
                else if (cmd.equals("freq")) {
                    Map<Character,Double> baseFrequencies = instance.getBaseCountFrequency(allSymbols);
                    output = instance.baseFrequencyToString(baseFrequencies);
                }
                else { // CG content
                    int gc = instance.gcContent();
                    double gcpercent = (double)gc/(double)instance.baseCount;
                    output = "total bases: " + instance.baseCount + "\n\n" + "\tCG Content = " + gc + "\n\tGC Percentage = " + String.format("%.4f", gcpercent) +"\n";
                }
                System.out.print(output);
            }
        }
        catch( ParseException exp ) {
            // oops, something went wrong
            LOG.exiting( "Parsing failed.  Reason: ", exp.getMessage() );
            System.exit(1);
        }


    }

}
